package com.tezign.intelligence.mini.dam.aifamilyconsultant.util;

import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationParam;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.aigc.generation.models.QwenParam;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import com.alibaba.dashscope.exception.NoApiKeyException;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import com.squareup.okhttp.OkHttpClient;
import com.squareup.okhttp.Request;
import com.squareup.okhttp.RequestBody;
import com.squareup.okhttp.Response;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionChunk;
import com.volcengine.ark.runtime.model.completion.chat.ChatCompletionRequest;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessage;
import com.volcengine.ark.runtime.model.completion.chat.ChatMessageRole;
import com.volcengine.ark.runtime.service.ArkService;
import io.github.cdimascio.dotenv.Dotenv;
import io.reactivex.BackpressureStrategy;
import io.reactivex.Flowable;
import lombok.extern.slf4j.Slf4j;
import okhttp3.MediaType;

import java.util.ArrayList;
import java.util.Arrays;
import java.util.List;
import java.util.StringJoiner;

import static com.tezign.intelligence.mini.dam.aifamilyconsultant.constant.Constant.DASHSCOPE_API_KEY;
import static com.tezign.intelligence.mini.dam.aifamilyconsultant.constant.Constant.VOLCENGINE_API_KEY;

/**
 * todo
 *
 * @Author ouyangkang
 * @date 2024/11/18
 */

@Slf4j
public class ChatUtil {

    public static final String QWEN_TURBO = "qwen-turbo-1101";
    public static final String MODEL = "doubao-pro-256k-241115";


    private static final String API_BASE = "https://api.deepseek.com/v1";
    private final String apiKey;
    private final OkHttpClient client;

    public ChatUtil(String apiKey) {
        this.apiKey = apiKey;
        this.client = new OkHttpClient();
    }

    public static Flowable<GenerationResult> streamCall(List<Message> userMsgs) throws Exception {
        Dotenv dotenv = Dotenv.load();
        String dashscopeApiKey = dotenv.get(DASHSCOPE_API_KEY);
        if (dashscopeApiKey == null) {
            throw new NoApiKeyException();
        }
        GenerationParam param = GenerationParam.builder()
                .apiKey(dashscopeApiKey)
                .model(QWEN_TURBO)
                .messages(userMsgs)
                .resultFormat(GenerationParam.ResultFormat.MESSAGE)
                .topP(0.8)
                .incrementalOutput(true)
                .build();

        Generation gen = new Generation();
        Flowable<GenerationResult> result = gen.streamCall(param);
        return result;
    }

    public static String callWithMessage(String role, String content) throws Exception {
        log.info("role={} content={}", role, content);
        StringJoiner stringJoiner = new StringJoiner("");
        Dotenv dotenv = Dotenv.load();
        String apiKey = dotenv.get(VOLCENGINE_API_KEY);
        ArkService service = ArkService.builder().apiKey(apiKey).build();
        final List<ChatMessage> messages = new ArrayList<>();
        final ChatMessage systemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content(role).build();
        final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content(content).build();
        messages.add(systemMessage);
        messages.add(userMessage);

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(MODEL)
                .messages(messages)
                .build();

        service.createChatCompletion(chatCompletionRequest).getChoices().forEach(choice -> stringJoiner.add(choice.getMessage().getContent().toString()));
        service.shutdownExecutor();
        return stringJoiner.toString();
    }

    public static void main(String[] args) throws Exception {
		String string = callWithMessage("你是一个程序员","请编写一个Python函数 find_prime_numbers，该函数接受一个整数 n 作为参数，并返回一个包含所有小于 n 的质数（素数）的列表。质数是指仅能被1和其自身整除的正整数，如2, 3, 5, 7等。");
		System.out.println(string);
//		Dotenv dotenv = Dotenv.load();
//		String apiKey = dotenv.get(VOLCENGINE_API_KEY);
//		ArkService service = ArkService.builder().apiKey(apiKey).build();
//		System.out.println("\n----- standard request -----");
//		final List<ChatMessage> messages = new ArrayList<>();
//		final ChatMessage systemMessage = ChatMessage.builder().role(ChatMessageRole.SYSTEM).content("你是豆包，是由字节跳动开发的 AI 人工智能助手").build();
//		final ChatMessage userMessage = ChatMessage.builder().role(ChatMessageRole.USER).content("常见的十字花科植物有哪些？").build();
//		messages.add(systemMessage);
//		messages.add(userMessage);
//
//		ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
//				.model(MODEL)
//				.messages(messages)
//				.build();
//
//		service.createChatCompletion(chatCompletionRequest).getChoices().forEach(choice -> System.out.println(choice.getMessage().getContent()));
//
//		// shutdown service
//		service.shutdownExecutor();
    }

    public static Flowable<ChatCompletionChunk> doubaoChat(List<ChatMessage> message) throws Exception {
        Dotenv dotenv = Dotenv.load();
        String apiKey = dotenv.get(VOLCENGINE_API_KEY);
        ArkService service = ArkService.builder().apiKey(apiKey).build();
        ChatCompletionRequest streamChatCompletionRequest = ChatCompletionRequest.builder()
                .model(MODEL)
                .messages(message)
                .build();
        return service.streamChatCompletion(streamChatCompletionRequest);
    }

//    public static Flowable<ChatCompletionChunk> deepseekChat(List<ChatMessage> messages) throws Exception {
//        Dotenv dotenv = Dotenv.load();
//        String apiKey = dotenv.get(VOLCENGINE_API_KEY);
//        DeepSeekClient deepSeekClient = new DeepSeekClient(apiKey);
//        return deepSeekClient.streamChatCompletion(messages);
//    }
}
